// extractpdf.cpp
#include "extractpdf.h"
#include <poppler-config.h>

#include <GlobalParams.h>
#include <PDFDoc.h>
#include <PDFDocFactory.h>
#include "ImageOutputDev.h"
#include <memory>
#include <stdexcept>

class CImageOutputDev : public ImageOutputDev {
public:
    CImageOutputDev(char* fileRootA, GBool pageNamesA, GBool listImagesA)
        : ImageOutputDev(fileRootA, pageNamesA, listImagesA) {
    }
    void setImageCallBack(PageDataCallback callback, int pageNumber) {
        callback_ = callback;
        pageNumber_ = pageNumber;
    }
    void writeRawImage(Stream* str, const char* ext) override;

private:
    std::vector<char> buffer_;
    PageDataCallback callback_;
    int pageNumber_;
};


void CImageOutputDev::writeRawImage(Stream* str, const char* ext) {
    if (callback_) {
        str = str->getNextStream();
        str->reset();
        buffer_.clear();
        int c;
        while ((c = str->getChar()) != EOF) {
            buffer_.push_back(static_cast<char>(c));
        }
        str->close();

        if (!callback_(pageNumber_, ContentType::content_img, buffer_, ext)) {
            // If callback returns false, we should stop processing
           // throw std::runtime_error("Extraction stopped by callback");
        }
    }
    else {
        ImageOutputDev::writeRawImage(str, ext);
    }
}

PdfExtractor::PdfExtractor() : m_numPages(0), m_imgOut(nullptr) {
    globalParams = new GlobalParams();
}

PdfExtractor::~PdfExtractor() = default;

bool PdfExtractor::setupFilePath(const std::string& pdfPath) {
    globalParams = new GlobalParams();

    GooString* ownerPW = nullptr;
    GooString* userPW = nullptr;

    GooString gooPdfPath(pdfPath.c_str());
    m_doc.reset(PDFDocFactory().createPDFDoc(gooPdfPath, ownerPW, userPW));

    if (!m_doc || !m_doc->isOk()) {
        return false;
    }

    m_popplerDoc.reset(poppler::document::load_from_file(pdfPath));
    if (!m_popplerDoc || m_popplerDoc->is_locked()) {
        return false;
    }
    m_numPages = m_popplerDoc->pages();
    return true;
}

bool PdfExtractor::extractPdfTextAndImagesFromPage(int pageIdx, PageDataCallback callback, ImageFormat imageFormat, ExtractionMode mode) {
    if (pageIdx >= m_numPages || pageIdx < 0) {
        return false;
    }

    std::unique_ptr<poppler::page> page(m_popplerDoc->create_page(pageIdx));
    if (page) {
        if (mode == ExtractionMode::TextOnly || mode == ExtractionMode::TextAndImages) {
            poppler::byte_array text = page->text().to_utf8();
            std::vector<char> textVector(text.begin(), text.end());
            if (!callback(pageIdx + 1, ContentType::content_text, textVector, "txt")) {
                return false;
            }
        }
        if (mode == ExtractionMode::ImagesOnly || mode == ExtractionMode::TextAndImages) {
            int firstPage = pageIdx + 1;
            int lastPage = firstPage + 1;
            bool listImages = false;
            std::string outDir = "./out/";

#ifdef _WIN32
            CreateDirectoryA(outDir.c_str(), NULL);
#else
            mkdir(outDir.c_str(), 0777);
#endif
            CImageOutputDev* imgOut = new CImageOutputDev(const_cast<char*>(outDir.c_str()), gTrue, listImages);
            if ( m_imgOut == nullptr ) {
                m_imgOut = imgOut;
            }
            m_imgOut->enablePNG(gTrue);
            m_imgOut->enableTiff(gTrue);
            m_imgOut->enableJpeg(gTrue);
            m_imgOut->enableJpeg2000(gTrue);
            m_imgOut->enableJBig2(gTrue);
            m_imgOut->enableCCITT(gTrue);
            m_imgOut->setImageCallBack(callback, pageIdx + 1);
            m_doc->displayPages(m_imgOut, firstPage, lastPage, 72, 72, 0, gTrue, gFalse, gFalse);
            delete m_imgOut;
            m_imgOut = nullptr;
        }
    }
    else {
        return false;
    }

    return true;
}

int PdfExtractor::getPageCount() const {
    return m_numPages;
}

bool PdfExtractor::extractAllPdfTextAndImages(const std::string& pdfPath, PageDataCallback callback, ImageFormat imageFormat, ExtractionMode mode) {
    // Create a PDF document object using Poppler's PDFDocFactory
    globalParams = new GlobalParams();

    GooString* ownerPW = nullptr;
    GooString* userPW = nullptr;

    // Convert std::string to GooString
    GooString gooPdfPath(pdfPath.c_str());

    // Call createPDFDoc with the correct type
    std::unique_ptr<PDFDoc> doc(PDFDocFactory().createPDFDoc(gooPdfPath, ownerPW, userPW));

    // Clean up password strings if they were used
    delete userPW;
    delete ownerPW;

    // Check if the document is valid
    if (!doc || !doc->isOk())
    {
        return false; // Exit function if document is not OK
    }

    // Load the document using Poppler's document class
    std::unique_ptr<poppler::document> popplerDoc(poppler::document::load_from_file(pdfPath));
    if (!popplerDoc || popplerDoc->is_locked())
    {
        return false; // Exit function if document is locked or failed to load
    }

    int numPages = popplerDoc->pages();
    for (int i = 0; i < numPages; ++i)
    {
        std::unique_ptr<poppler::page> page(popplerDoc->create_page(i));
        if (page)
        {
            if (mode == ExtractionMode::TextOnly || mode == ExtractionMode::TextAndImages)
            {
                poppler::byte_array text = page->text().to_utf8();
                std::vector<char> textVector(text.begin(), text.end());
                if (!callback(i + 1, ContentType::content_text, textVector, "txt")) {
                    return false; // Stop processing if callback returns false
                }
            }

            if (mode == ExtractionMode::ImagesOnly || mode == ExtractionMode::TextAndImages)
            {
                int firstPage = i + 1;
                int lastPage = firstPage + 1;
                bool listImages = false;
                // Create output directory path
                std::string outDir = "./out/";
#ifdef _WIN32
                CreateDirectoryA(outDir.c_str(), NULL);
#else
                mkdir(outDir.c_str(), 0777);
#endif
                CImageOutputDev* imgOut = new CImageOutputDev(const_cast<char*>(outDir.c_str()), gTrue, listImages);
                imgOut->enablePNG(gTrue);
                imgOut->enableTiff(gTrue);
                imgOut->enableJpeg(gTrue);
                imgOut->enableJpeg2000(gTrue);
                imgOut->enableJBig2(gTrue);
                imgOut->enableCCITT(gTrue);
                if (callback != nullptr) {
                    imgOut->setImageCallBack(callback, firstPage);
                }
                doc->displayPages(imgOut, firstPage, lastPage, 72, 72, 0, gTrue, gFalse, gFalse);
                delete imgOut;
            }
        }
        else
        {
            // Handle error if page creation fails
            return false;
        }
    }
    return true;
}
